# - 创建函数 fn，函数包含参数 a, b，计算，函数内校验 a b 是否可以相乘，并返回 a b 矩阵相乘的结果
# - 若参数 a 是一维数据，则在扩展一个维度到形状的最前面
# - 若参数 a 或 b不是二维矩阵，要给出提示
# - 若参数 a 和 b都是二维矩阵，但是不能相乘 要给出提示

# - 创建一组数据 x，数据特征数为 2，共 10 条数据
# - x 的所有元素取值范围是 0~1
# - 创建一组数据 w，形状为 (2, 5)
# - 调用 fn(x, w)

# - 重新创建 x，这次 x 形状为 (2,)
# - 调用 fn(x, w)

# - 重新创建 w，这次 w 形状为 (5, 2)
# - 调用 fn(x, w)

import torch


# def fn(a, b):
#     if a.dim() == 1:
#         a = a.unsqueeze(0)
#         print(f"a 是一维数据，已扩展维度，新形状: {a.shape}")
#
#     if a.dim() != 2 or b.dim() != 2:
#         raise ValueError(f"参数必须是二维矩阵，a的维度: {a.dim()}, b的维度: {b.dim()}")
#
#     if a.shape[1] != b.shape[0]:
#         raise ValueError(f"矩阵无法相乘，a的形状: {a.shape}, b的形状: {b.shape}。需要a的列数({a.shape[1]})等于b的行数({b.shape[0]})")

def fn(a, b):
    if len(a.shape) == 1:
        a = a.unsqueeze(0)
    '''
    python中关于错误的调试可以使用断言机制assert 和 异常机制try-except-raise
    断言机制：仅用于调试阶段，验证程序运行到此处必须满足的条件，多用于函数参数的判断
    语法格式：assert 条件, "提示语句 AssertionError"    （条件不满足自动触发AssertionError）
    '''
    assert len(a.shape) == 2 and len(b.shape) == 2, "a或b不是矩阵"
    assert a.shape[1] == b.shape[0], "a和b的维度不匹配，无法进行矩阵相乘"

    return a @ b


# 测试数据1
x = torch.rand(10, 2)
w = torch.rand(2, 5)
print(fn(x, w).shape)

# 测试数据2
x = torch.rand(2)
print(fn(x, w).shape)

# 测试数据3
w = torch.rand(5, 2)
# print(fn(x, w).shape) AssertionError: a和b的维度不匹配，无法进行矩阵相乘
# print(torch.mm(x, w).shape) RuntimeError: self must be a matrix
# print(torch.matmul(x, w).shape) RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x2 and 5x2)

# 测试数据4
x = torch.rand(10, 2)
w = torch.rand(5)
# print(torch.matmul(x, w).shape)  RuntimeError: size mismatch, got input (10), mat (10x2), vec (5)
